﻿using System;
using System.Threading;
using System.Threading.Tasks;

namespace Prometheus.Tests;

/// <summary>
/// A delayer that seems to work as usual except it can be instructed to end all waits immediately.
/// </summary>
/// <remarks>
/// Thread-safe.
/// </remarks>
public sealed class BreakableDelayer : IDelayer
{
    /// <summary>
    /// Ends all delays and pretends the timers all elapsed.
    /// </summary>
    public void BreakAllDelays()
    {
        CancellationTokenSource old;

        lock (_lock)
        {
            // Have to replace CTS first to ensure that any new calls get new CTS.
            // Very important because canceling the CTS actually executes code until the next await.
            old = _cts;
            _cts = new CancellationTokenSource();
        }

        old.Cancel();
        old.Dispose();
    }

    private CancellationTokenSource _cts = new CancellationTokenSource();
    private readonly object _lock = new object();

    public async Task Delay(TimeSpan duration)
    {
        CancellationToken cancel;

        lock (_lock)
            cancel = _cts.Token;

        try
        {
            await Task.Delay(duration, cancel);
        }
        catch (TaskCanceledException)
        {
        }
    }

    public async Task Delay(TimeSpan duration, CancellationToken requestedCancel)
    {
        CancellationTokenSource callCts;
        CancellationToken cancel;

        lock (_lock)
        {
            callCts = CancellationTokenSource.CreateLinkedTokenSource(_cts.Token, requestedCancel);
            cancel = callCts.Token;
        }

        try
        {
            await Task.Delay(duration, cancel);
        }
        catch (TaskCanceledException)
        {
            requestedCancel.ThrowIfCancellationRequested();
        }
        finally
        {
            callCts.Dispose();
        }
    }
}
